import os
import numpy as np
import matplotlib.pyplot as plt
from torchvision import transforms
from torch.utils.data import Dataset
from PIL import Image
from utils import *
from glob import glob
import albumentations as A
import random
class Neu_Seg_Competition_Dataset_nation(Dataset):
    def __init__(self, data_path, data_type,transform=None,seed=24561):  #B级涨点神秘数字
        self.data_path = data_path
        self.data_type = data_type
        self.image_names_path = glob(os.path.join(data_path, 'Img','*.jpg'))
        self.to_tensor = transforms.ToTensor()
        # self.mask_names = glob(os.path.join(data_path, 'annotations',data_type,'*.png'))
        self.transform_ = transform
        np.random.seed(seed)
        np.random.shuffle(self.image_names_path)
        split_idx = int(0.4 * len(self.image_names_path))
        if data_type == 'training':
            self.image_names_path = self.image_names_path[:split_idx]
        else:
            self.image_names_path = self.image_names_path[split_idx:]
    def __len__(self):
        return len(self.image_names_path)
    def __getitem__(self, idx):
        image_name = self.image_names_path[idx][-10:-4]
        image = Image.open(os.path.join(self.data_path,'Img', f'{image_name}.jpg'))
        mask = Image.open(os.path.join(self.data_path,'Lab', f'{image_name}.png'))
        if self.transform_:
            # 转换为NumPy数组以便于Albumentations处理
            image_np = np.array(image)
            mask_np = np.array(mask)
            aug = self.transform_(image=image_np, mask=mask_np)
            # aug = self.transform_(image=(image), mask=(mask))
            img = Image.fromarray(aug['image'])  # Convert back to PIL
            # mask = aug['mask']#numpy array
            # return self.to_tensor(img), torch.from_numpy(mask), image_name
            mask = Image.fromarray(aug['mask'])
            return self.to_tensor(img), PIL_Image_ToTensor(mask), image_name


        return self.to_tensor(image), PIL_Image_ToTensor(mask),image_name

if __name__ == '__main__':
    import torch
    # t_train = None
    import torch
    import albumentations as A
    import matplotlib.pyplot as plt
    t_train=None
    t_train = A.Compose([
        A.RandomBrightnessContrast(brightness_limit=0.5, contrast_limit=0.5, p=1.0), #亮度、对比度

    ])
    train_dataset = Neu_Seg_Competition_Dataset_nation(data_path='.', data_type='test', transform=t_train)
    print(f'Total samples in dataset: {len(train_dataset)}')

    # 获取增强后的样本
    sample_idx = 0  # You can change this index to visualize different samples
    img_tensor, mask_tensor, image_name = train_dataset[sample_idx]

    # Convert tensors back to numpy for visualization
    img = img_tensor.permute(1, 2, 0).numpy()  # Change to HWC format
    mask = mask_tensor.numpy()  # Assuming mask is a single channel

    plt.figure(figsize=(12, 6))

    # Original image
    plt.subplot(1, 2, 1)
    plt.imshow(img, cmap='gray')
    plt.title(f'Augmented Image: {image_name}')
    plt.axis('off')

    # Mask
    plt.subplot(1, 2, 2)
    plt.imshow(mask, cmap='gray')
    plt.title('Augmented Mask')
    plt.axis('off')

    plt.show()


    # print(torch.unique(train_dataset[120][1]))
    # fig,axs = plt.subplots(1,2,figsize=(10,40))
    # # axs[0,0] = plt.imshow(train_dataset[120][0])
    # axs[0,1] = plt.imshow(train_dataset[120][1])
    # plt.show()